#include "kernel_operator.h"
#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

template <typename T>
struct SetValueHelper
{
    __aicore__ static void Set(AscendC::GlobalTensor<T> &z, uint32_t i, T value)
    {
        z.SetValue(i, value);
    }
};

template <>
struct SetValueHelper<double>
{
    __aicore__ static void Set(AscendC::GlobalTensor<double> &z, uint32_t i, double value)
    {
        z.SetValue(i, value);
    }
};

class KernelEye
{
public:
    __aicore__ inline KernelEye() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z,
                                uint32_t size,
                                uint32_t tileNum,
                                uint32_t num_rows,
                                uint32_t num_columns,
                                uint32_t batch_shape_numel,
                                uint32_t dtype)
    {
        uint32_t blockNum = AscendC::GetBlockNum();
        uint32_t blockIdx = AscendC::GetBlockIdx();

        // Calculate base block size and remainder for even distribution
        uint32_t baseBlockSize = size / blockNum;
        uint32_t remainder = size % blockNum;

        // Calculate alignment requirements
        uint32_t alignmentInBytes = 64;
        uint32_t elementsPerAlignment = alignmentInBytes / sizeof(DTYPE_Y);
        // Ensure elementsPerAlignment is at least 1
        elementsPerAlignment = MAX(elementsPerAlignment, 1);
        
        // Number of aligned chunks in the data
        uint32_t totalAlignedChunks = (size + elementsPerAlignment - 1) / elementsPerAlignment;
        
        // Distribute aligned chunks to blocks
        uint32_t chunksPerBlock = totalAlignedChunks / blockNum;
        uint32_t remainderChunks = totalAlignedChunks % blockNum;
        
        // Calculate start and length (in aligned chunks)
        uint32_t startChunk = blockIdx * chunksPerBlock + MIN(blockIdx, remainderChunks);
        uint32_t numChunks = chunksPerBlock + (blockIdx < remainderChunks ? 1 : 0);
        
        // Convert to elements
        this->startPos = startChunk * elementsPerAlignment;
        this->blockLength = MIN(numChunks * elementsPerAlignment, size - this->startPos);

        this->num_rows = num_rows;
        this->num_columns = num_columns;
        this->batch_shape_numel = batch_shape_numel;

        AscendC::printf("blockLength: %d, startPos: %d, num_rows: %d, num_columns: %d, batch_shape_numel: %d\n", this->blockLength, this->startPos, this->num_rows, this->num_columns, this->batch_shape_numel);

        zGm.SetGlobalBuffer((__gm__ DTYPE_Y *)z + this->startPos, this->blockLength);
    }
    __aicore__ inline void Process()
    {
        for (uint32_t i = 0; i < this->blockLength; i++)
        {
            uint32_t flat_idx = i + this->startPos;
            uint32_t matrix_size = num_rows * num_columns;

            // Which matrix in the batch we're processing
            uint32_t batch_idx = flat_idx / matrix_size;

            // Index within the current matrix
            uint32_t elem_idx = flat_idx % matrix_size;

            uint32_t row = elem_idx / num_columns;
            uint32_t col = elem_idx % num_columns;

            if (row == col)
            {
                SetValueHelper<DTYPE_Y>::Set(zGm, i, 1);
            }
            else
            {
                SetValueHelper<DTYPE_Y>::Set(zGm, i, 0);
            }
        }
    }

private:
    AscendC::GlobalTensor<DTYPE_Y> zGm;
    uint32_t blockLength;
    uint32_t num_rows;
    uint32_t num_columns;
    uint32_t batch_shape_numel;
    uint32_t dtype;
    uint32_t startPos;
};

extern "C" __global__ __aicore__ void eye(GM_ADDR y, GM_ADDR y_ref, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelEye op;
    op.Init(y, y_ref,
            tiling_data.size,
            tiling_data.tileNum,
            tiling_data.num_rows,
            tiling_data.num_columns,
            tiling_data.batch_shape_numel,
            tiling_data.dtype);
    op.Process();
}